package com.gitee.mp.core.connector;

import com.alibaba.fastjson.JSONObject;
import com.gitee.mp.configure.MilvusConnection;
import com.gitee.mp.infra.constants.MilvusPlusConstants;
import com.gitee.mp.infra.struct.LazyLoader;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.response.QueryResp;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.CollectionUtils;

import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * @author hongda.li
 */
@Slf4j
@Getter
public class MilvusConnector {

    private final String connectorId;

    private final MilvusConnection connection;

    private final LazyLoader<MilvusClientV2> client;

    public MilvusConnector(MilvusConnection connection) {
        this.connection = connection;
        this.connectorId = connection.getConnectorId();
        this.client = new LazyLoader<>(this::buildClientByConnection);
    }

    @SuppressWarnings("ExtractMethodRecommender")
    private MilvusClientV2 buildClientByConnection() {
        ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder();

        if (this.connection.getUrl() != null) {
            builder.uri(this.connection.getUrl());
        }
        if (this.connection.getUsername() != null) {
            builder.username(this.connection.getUsername());
        }
        if (this.connection.getPassword() != null) {
            builder.password(this.connection.getPassword());
        }
        if (this.connection.getIdleTimeout() != null) {
            builder.idleTimeoutMs(this.connection.getIdleTimeout().toMillis());
        }
        if (this.connection.getConnectTimeout() != null) {
            builder.connectTimeoutMs(this.connection.getConnectTimeout().toMillis());
        }
        if (this.connection.getKeepAliveTime() != null) {
            builder.keepAliveTimeMs(this.connection.getKeepAliveTime().toMillis());
        }
        if (this.connection.getKeepAliveTimeout() != null) {
            builder.keepAliveTimeoutMs(this.connection.getKeepAliveTimeout().toMillis());
        }

        return new MilvusClientV2(builder.build());
    }

    /**
     * 插入数据
     *
     * @param collectionName 集合名称
     * @param dataList       字段名称以及字段对应的数据
     * @return 插入数量
     */
    public long insert(String collectionName, List<Map<String, Object>> dataList) {
        return this.client.load()
                .insert(InsertReq.builder()
                        .collectionName(collectionName)
                        .data(dataList.stream().map(JSONObject::new).toList())
                        .build())
                .getInsertCnt();
    }

    /**
     * 按主键更新数据
     * 数据不存在则新增
     *
     * @param collectionName 集合名称
     * @param dataList       字段名称以及字段对应的数据
     * @return 更新或插入数量
     */
    public long upsert(String collectionName, List<Map<String, Object>> dataList) {
        return this.client.load()
                .upsert(UpsertReq.builder()
                        .collectionName(collectionName)
                        .data(dataList.stream().map(JSONObject::new).toList())
                        .build())
                .getUpsertCnt();
    }

    /**
     * 根据表达式删除数据
     *
     * @param collectionName 集合名称
     * @param expression     筛选的表达式
     * @return 删除的数量
     */
    public long delete(String collectionName,
                       String expression,
                       List<Object> primaryKeyList) {
        DeleteReq.DeleteReqBuilder<?, ?> builder = DeleteReq.builder().collectionName(collectionName);

        if (CollectionUtils.isEmpty(primaryKeyList)) {
            // 根据表达式删除
            builder.filter(expression);
        } else {
            // 根据主键删除
            builder.ids(primaryKeyList);
        }

        return this.client.load().delete(builder.build()).getDeleteCnt();
    }

    /**
     * 查询数据
     *
     * @param collectionName 集合名称
     * @param expression     表达式
     * @param limit          查询数量
     * @param offset         偏移量
     * @param outputFields   查询字段
     * @return 查询结果
     */
    public List<Map<String, Object>> query(String collectionName,
                                           String expression,
                                           List<Object> primaryKeyList,
                                           Long limit,
                                           Long offset,
                                           List<String> outputFields) {
        QueryReq.QueryReqBuilder<?, ?> builder = QueryReq.builder().collectionName(collectionName);

        if (limit != null) {
            builder.limit(limit);
        }
        if (offset != null) {
            builder.offset(offset);
        }
        if (!CollectionUtils.isEmpty(outputFields)) {
            builder.outputFields(outputFields);
        }

        if (CollectionUtils.isEmpty(primaryKeyList)) {
            // 根据表达式查询
            builder.filter(expression);
        } else {
            // 根据主键查询
            builder.ids(primaryKeyList);
        }

        return this.client.load()
                .query(builder.build())
                .getQueryResults()
                .stream()
                .map(QueryResp.QueryResult::getEntity)
                .toList();
    }

    /**
     * 检索向量
     *
     * @param collectionName 集合名称
     * @param vectorList     向量集合
     * @param expression     过滤表达式
     * @param limit          检索数量
     * @param offset         偏移量
     * @param outputFields   查询字段
     * @return 检索结果
     */
    public List<Map<String, Object>> search(String collectionName,
                                            List<List<?>> vectorList,
                                            String expression,
                                            Long limit,
                                            Long offset,
                                            List<String> outputFields) {
        SearchReq.SearchReqBuilder<?, ?> builder = SearchReq.builder().collectionName(collectionName);

        if (limit != null) {
            builder.topK(limit.intValue());
        }
        if (offset != null) {
            builder.offset(offset);
        }
        if (!CollectionUtils.isEmpty(outputFields)) {
            builder.outputFields(outputFields);
        }

        return this.client.load()
                .search(builder.data(vectorList).filter(expression).build())
                .getSearchResults()
                .iterator()
                .next()
                .stream()
                .map(res -> {
                    Map<String, Object> result = new LinkedHashMap<>(res.getEntity());
                    result.put(MilvusPlusConstants.FIELD_ID_NAME, res.getId());
                    result.put(MilvusPlusConstants.FIELD_DISTANCE_NAME, res.getDistance());
                    return result;
                })
                .toList();
    }

    /**
     * 移除客户端
     */
    public boolean remove() {
        try {
            this.client.load().close(TimeUnit.MINUTES.toSeconds(1));
            this.client.unload();
            return true;
        } catch (InterruptedException e) {
            log.error("Error in close milvus", e);
            return false;
        }
    }
}
